import httpx
import json
import os
from typing import Dict, Optional, Any, AsyncGenerator
from core.config import settings
import aioredis
import logging
from services.chat_service import ChatService

# 设置日志记录器
logger = logging.getLogger(__name__)

class MaxKBAgent:
    def __init__(self):
        # 从配置中获取API信息 - 使用OpenAI兼容的URL格式
        self.base_url = os.environ.get("MAXKB_BASE_URL") or settings.MAXKB_BASE_URL
        self.api_key = os.environ.get("MAXKB_API_KEY") or settings.MAXKB_API_KEY
        self.chat_id = os.environ.get("MAXKB_Chat_ID") or settings.MAXKB_Chat_ID
        # 使用Bearer token格式的认证头部
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        # Redis配置 - 从环境变量或设置中获取
        self.redis_config = {
            "host": os.environ.get("REDIS_HOST") or settings.REDIS_HOST,
            "port": int(os.environ.get("REDIS_PORT") or settings.REDIS_PORT),
            "db": int(os.environ.get("REDIS_DB") or settings.REDIS_DB),
            "password": os.environ.get("REDIS_PASSWORD") or settings.REDIS_PASSWORD
        }
        # 如果密码为空，删除该键
        if self.redis_config.get("password") in ("", None):
            self.redis_config.pop("password")
        self.redis_client = None
        self.cache_ttl = 3600  # 缓存过期时间，单位秒（1小时）
        # 初始化ChatService用于存储聊天记录
        self.chat_service = ChatService()
    
    def _get_completions_url(self) -> str:
        """
        获取OpenAI兼容的chat completions URL
        
        Returns:
            完整的chat completions URL
        """
        # 确保基础URL末尾没有斜杠
        base_url = self.base_url.rstrip('/')
        # 根据MaxKB官方文档，使用OpenAI兼容格式: Base URL/chat/completions
        return f"{base_url}/chat/completions"
    
    async def get_redis_client(self):
        """
        获取Redis客户端连接
        
        Returns:
            Redis客户端实例
        """
        if not self.redis_client:
            try:
                # 构建Redis连接地址
                redis_address = (self.redis_config["host"], self.redis_config["port"])
                
                # 构建连接参数
                conn_kwargs = {}
                if "password" in self.redis_config:
                    conn_kwargs["password"] = self.redis_config["password"]
                conn_kwargs["db"] = self.redis_config["db"]
                
                # 使用aioredis 1.3.1的连接方式
                self.redis_client = await aioredis.create_redis_pool(
                    redis_address,
                    encoding="utf-8",
                    **conn_kwargs
                )
                # 测试连接
                await self.redis_client.ping()
                logger.info("Redis连接成功")
            except Exception as e:
                logger.error(f"Redis连接失败: {str(e)}")
                self.redis_client = None
        return self.redis_client
    
    def _get_cache_key(self, query: str, chat_id: str, account_no: str) -> str:
        """
        构建缓存键
        
        Args:
            query: 用户查询内容
            chat_id: 对话ID
            account_no: 用户账号标识
            
        Returns:
            缓存键字符串
        """
        return f"maxkb_chat:{account_no}:{chat_id}:{hash(query)}"
    
    async def get_cached_response(self, query: str, chat_id: str, account_no: str = "") -> Optional[Dict[str, Any]]:
        """
        从缓存中获取响应
        
        Args:
            query: 用户查询内容
            chat_id: 对话ID
            account_no: 用户账号标识
            
        Returns:
            缓存的响应，如果没有缓存则返回None
        """
        try:
            redis_client = await self.get_redis_client()
            if not redis_client:
                return None
            
            cache_key = self._get_cache_key(query, chat_id, account_no)
            cached_data = await redis_client.get(cache_key)
            
            if cached_data:
                logger.info(f"从缓存获取响应: {cache_key}, account_no={account_no}")
                return json.loads(cached_data)
            return None
        except Exception as e:
            logger.error(f"获取缓存失败: {str(e)}, account_no={account_no}")
            return None
    
    async def set_cached_response(self, query: str, chat_id: str, account_no: str = "", response: Dict[str, Any] = None):
        """
        设置缓存响应
        
        Args:
            query: 用户查询内容
            chat_id: 对话ID
            account_no: 用户账号标识
            response: 要缓存的响应数据
        """
        try:
            redis_client = await self.get_redis_client()
            if not redis_client or not response:
                return
            
            cache_key = self._get_cache_key(query, chat_id, account_no)
            await redis_client.setex(
                cache_key,
                self.cache_ttl,
                json.dumps(response)
            )
            logger.info(f"设置缓存响应: {cache_key}, TTL: {self.cache_ttl}秒, account_no={account_no}")
        except Exception as e:
            logger.error(f"设置缓存失败: {str(e)}, account_no={account_no}")
    
    async def get_application_profile(self) -> Dict[str, Any]:
        """
        获取应用相关信息
        
        Returns:
            应用信息
        """
        try:
            endpoint = "application/profile"
            full_url = self._construct_url(endpoint)
            
            async with httpx.AsyncClient(timeout=30.0) as client:
                response = await client.get(full_url, headers=self.headers)
                response.raise_for_status()
                return response.json()
        except httpx.HTTPError as e:
            return {"error": f"HTTP错误: {str(e)}", "status_code": response.status_code if 'response' in locals() else None}
        except Exception as e:
            return {"error": f"发生错误: {str(e)}"}
    
    async def send_message(self, message: str, chat_id: Optional[str] = None, re_chat: bool = False, account_no: str = "") -> Dict[str, Any]:
        """
        发送消息到指定对话
        
        Args:
            message: 发送的消息内容
            chat_id: 对话ID，默认为配置中的MAXKB_Chat_ID
            re_chat: 是否重新开始对话
            
        Returns:
            消息响应
        """
        try:
            # 使用配置中的对话ID或传入的对话ID
            target_chat_id = chat_id or self.chat_id
            if not target_chat_id:
                raise ValueError("对话ID不能为空")
            
            # 如果不是重新开始对话，检查缓存
            if not re_chat:
                cached_response = await self.get_cached_response(message, target_chat_id, account_no)
                if cached_response:
                    logger.info(f"使用缓存响应，查询: {message[:50]}..., account_no={account_no}")
                    # 确保返回的响应包含chat_id和account_no
                    if "chat_id" not in cached_response:
                        cached_response["chat_id"] = target_chat_id
                    if "account_no" not in cached_response:
                        cached_response["account_no"] = account_no
                    return cached_response
            
            logger.info(f"缓存未命中或需要重新对话，查询: {message[:50]}...")
            
            # 构建API端点 - 使用正确的路径格式
            endpoint = f"chat_message/{target_chat_id}"
            full_url = self._construct_url(endpoint)
            
            # 构建请求体 - 添加流式标志字段
            payload = {
                "message": message,
                "re_chat": re_chat,
                "stream": False  # 非流式标志
            }
            
            logger.debug(f"发送请求到: {full_url}")
            logger.debug(f"请求体: {payload}")
            
            async with httpx.AsyncClient(timeout=60.0) as client:
                response = await client.post(full_url, headers=self.headers, json=payload)
                response.raise_for_status()
                result = response.json()
                logger.debug(f"API响应: {result}")
                
                # 确保响应中包含对话ID
                if "chat_id" not in result:
                    result["chat_id"] = target_chat_id
                
                # 如果不是重新开始对话且响应成功，缓存结果
                if not re_chat and "error" not in result:
                    # 确保结果中包含account_no
                    result["account_no"] = account_no
                    await self.set_cached_response(message, target_chat_id, account_no, result)
                    
                    # 保存聊天记录
                    if account_no and "content" in result:
                        save_ok = await self.chat_service.save_chat_message(
                            account_no=account_no,
                            user_msg=message,
                            assistant_msg=result["content"]
                        )
                        if save_ok:
                            logger.info(f"[account_no={account_no}] 聊天记录已保存")
                        # 异步生成对话摘要
                        import asyncio
                        asyncio.create_task(self.chat_service.generate_summary(account_no=account_no))
                
                return result
        except httpx.HTTPError as e:
            error_info = {"error": f"HTTP错误: {str(e)}"}
            if 'response' in locals():
                error_info["status_code"] = response.status_code
                try:
                    error_details = response.json()
                    if isinstance(error_details, dict):
                        error_info.update(error_details)
                except:
                    error_info["response_text"] = response.text if 'response' in locals() else ""
            return error_info
        except Exception as e:
            return {"error": f"发生错误: {str(e)}"}
    
    async def stream_chat(self, message: str, chat_id: Optional[str] = None, re_chat: bool = False, account_no: str = "") -> AsyncGenerator[Any, None]:
        """
        流式对话方法 - 与router中调用的方法名保持一致
        
        Args:
            message: 发送的消息内容
            chat_id: 对话ID，默认为配置中的MAXKB_Chat_ID
            re_chat: 是否重新开始对话
            
        Yields:
            流式响应的每个部分
        """
        # 直接调用现有的send_message_stream方法
        async for chunk in self.send_message_stream(message, chat_id, re_chat):
            # 尝试解析为JSON对象返回，方便router处理
            try:
                if isinstance(chunk, bytes):
                    chunk_str = chunk.decode('utf-8', errors='ignore')
                    yield json.loads(chunk_str)
                else:
                    yield chunk
            except (json.JSONDecodeError, UnicodeDecodeError):
                # 如果解析失败，返回原始chunk
                yield chunk
    
    async def send_message_stream(self, message: str, chat_id: Optional[str] = None, re_chat: bool = False, account_no: str = "") -> AsyncGenerator[bytes, None]:
        """
        流式发送消息到MaxKB，使用OpenAI兼容的API格式
        
        Args:
            message: 发送的消息内容
            chat_id: 对话ID，默认为配置中的MAXKB_Chat_ID
            re_chat: 是否重新开始对话
            
        Yields:
            流式响应的每个部分（以UTF-8编码的字节格式）
        """
        try:
            # 使用配置中的对话ID或传入的对话ID
            target_chat_id = chat_id or self.chat_id
            if not target_chat_id:
                raise ValueError("对话ID不能为空")
            
            # 获取OpenAI兼容的URL
            full_url = self._get_completions_url()
            
            # 构建OpenAI兼容的请求体
            payload = {
                "model": "gpt-3.5-turbo",  # 固定模型名称
                "messages": [
                    {
                        "role": "user",
                        "content": message
                    }
                ],
                "stream": True
            }
            
            logger.debug(f"发送流式请求到: {full_url}")
            logger.debug(f"请求体: {payload}")
            
            # 用于收集完整响应内容以便缓存
            full_response_content = []
            
            # 使用流式模式发送请求
            async with httpx.AsyncClient(timeout=None) as client:
                async with client.stream('POST', full_url, headers=self.headers, json=payload) as response:
                    response.raise_for_status()
                    
                    # 逐块获取响应数据
                    async for chunk in response.aiter_bytes():
                        # 解析和处理响应块
                        try:
                            chunk_str = chunk.decode('utf-8', errors='ignore')
                            
                            # 处理SSE格式的数据
                            if chunk_str.startswith('data: '):
                                data_part = chunk_str[6:]
                                try:
                                    # 解析JSON数据
                                    data = json.loads(data_part)
                                    
                                    # 提取内容
                                    content = data.get('content', '') or data.get('choices', [{}])[0].get('delta', {}).get('content', '')
                                    
                                    if content:
                                        full_response_content.append(content)
                                        
                                        # 包装响应数据，确保包含必要的字段
                                    wrapped_data = {
                                        "content": content,
                                        "chat_id": target_chat_id,
                                        "account_no": account_no
                                    }
                                    # 添加0.1秒延迟，降低流式输出速度，避免前端处理压力
                                    import asyncio
                                    await asyncio.sleep(0.1)
                                    yield json.dumps(wrapped_data).encode('utf-8')
                                except json.JSONDecodeError as e:
                                    logger.error(f"解析JSON错误: {e}")
                                    logger.debug(f"原始数据: {data_part}")
                        except Exception as e:
                            logger.error(f"解析响应块失败: {str(e)}")
                            # 继续处理下一个块
                    
                    # 发送完成信号
                    if full_response_content:
                        completion_data = {
                            "content": '',  # 完成信号不包含内容
                            "chat_id": target_chat_id,
                            "account_no": account_no,
                            "finished": True
                        }
                        yield json.dumps(completion_data).encode('utf-8')
            
            # 如果不是重新开始对话且有有效内容，缓存结果
            if not re_chat and full_response_content:
                # 构建要缓存的完整响应
                cached_response = {
                    "content": ''.join(full_response_content),
                    "chat_id": target_chat_id,
                    "account_no": account_no
                }
                await self.set_cached_response(message, target_chat_id, account_no, cached_response)
                
                # 保存聊天记录
                if account_no and "content" in cached_response:
                    save_ok = await self.chat_service.save_chat_message(
                        account_no=account_no,
                        user_msg=message,
                        assistant_msg=cached_response["content"]
                    )
                    if save_ok:
                        logger.info(f"[account_no={account_no}] 流式聊天记录已保存")
                    # 异步生成对话摘要
                    import asyncio
                    asyncio.create_task(self.chat_service.generate_summary(account_no=account_no))
        except httpx.HTTPError as e:
            error_info = {"error": f"HTTP错误: {str(e)}"}
            if 'response' in locals():
                error_info["status_code"] = response.status_code
                try:
                    error_info["response_text"] = response.text
                except:
                    pass
            yield json.dumps(error_info).encode('utf-8')
        except Exception as e:
            logger.error(f"发生错误: {str(e)}")
            yield json.dumps({"error": f"发生错误: {str(e)}"}).encode('utf-8')


# 创建单例实例
maxkb_agent = MaxKBAgent()


# 便捷函数
def get_maxkb_agent() -> MaxKBAgent:
    """
    获取MaxKB代理实例
    
    Returns:
        MaxKB代理实例
    """
    return maxkb_agent


async def maxkb_chat(query: str, conversation_id: Optional[str] = None, re_chat: bool = False, account_no: str = "") -> Dict[str, Any]:
    """
    MaxKB对话便捷函数
    
    Args:
        query: 用户查询内容
        conversation_id: 对话ID，可选，默认为配置中的MAXKB_Chat_ID
        re_chat: 是否重新开始对话
        
    Returns:
        对话响应
    """
    agent = get_maxkb_agent()
    
    try:
        # 调用send_message方法，传入查询内容和对话ID
        response = await agent.send_message(
            message=query,
            chat_id=conversation_id,
            re_chat=re_chat,
            account_no=account_no
        )
        
        # 确保响应中包含对话ID
        if "error" not in response:
            response["chat_id"] = conversation_id or agent.chat_id
            
        return response
    except Exception as e:
        return {"error": str(e)}


async def maxkb_chat_stream(query: str, conversation_id: Optional[str] = None, re_chat: bool = False, account_no: str = "") -> AsyncGenerator[bytes, None]:
    """
    便捷函数：与MaxKB进行流式对话
    
    Args:
        query: 用户查询
        conversation_id: 对话ID，可选，默认为配置中的MAXKB_Chat_ID
        re_chat: 是否重新开始对话
        
    Yields:
        流式响应的每个部分（字节格式）
    """
    agent = get_maxkb_agent()
    current_chat_id = conversation_id or agent.chat_id
    
    try:
        # 如果不是重新开始对话，先检查缓存
            if not re_chat:
                cached_response = await agent.get_cached_response(query, current_chat_id, account_no)
                if cached_response and "error" not in cached_response:
                    logger.info(f"使用缓存响应流式查询: {query[:50]}..., account_no={account_no}")
                
                
                if "content" in cached_response:
                   
                    content = cached_response["content"]
                    chunk_size = 100  # 每次返回的字符数
                    
                    for i in range(0, len(content), chunk_size):
                        chunk = content[i:i+chunk_size]
                        wrapped_chunk = {
                            "content": chunk,
                            "chat_id": current_chat_id,
                            "account_no": account_no
                        }
                        yield json.dumps(wrapped_chunk).encode('utf-8')
                        
                        
                        import asyncio
                        await asyncio.sleep(0.1)
                    
                   
                    yield json.dumps({"finished": True, "chat_id": current_chat_id, "account_no": account_no}).encode('utf-8')
                    
                   
                    if account_no and cached_response and "content" in cached_response:
                        save_ok = await agent.chat_service.save_chat_message(
                            account_no=account_no,
                            user_msg=query,
                            assistant_msg=cached_response["content"]
                        )
                        if save_ok:
                            logger.info(f"[account_no={account_no}] 缓存响应聊天记录已保存")
                        # 异步生成对话摘要
                        asyncio.create_task(agent.chat_service.generate_summary(account_no=account_no))
                    return
        # 流式发送消息
            async for chunk in agent.send_message_stream(
                message=query,
                chat_id=current_chat_id,
                re_chat=re_chat,
                account_no=account_no
            ):
            # 尝试解析SSE格式的响应
                chunk_str = chunk.decode('utf-8', errors='ignore').strip()
            
            # 处理SSE格式的数据块
            if chunk_str.startswith('data: '):
                data_part = chunk_str[6:]
                try:
                    
                    data = json.loads(data_part)
                    o
                    if "chat_id" not in data:
                        data["chat_id"] = current_chat_id
                    if "account_no" not in data:
                        data["account_no"] = account_no
                    # 重新编码为JSON并返回
                    yield json.dumps(data).encode('utf-8')
                except json.JSONDecodeError:
                    # 如果不是有效的JSON，原样返回
                    yield chunk
            else:
                # 非SSE格式数据，尝试直接解析为JSON
                try:
                    data = json.loads(chunk_str)
                    if "chat_id" not in data:
                        data["chat_id"] = current_chat_id
                    if "account_no" not in data:
                        data["account_no"] = account_no
                    yield json.dumps(data).encode('utf-8')
                    
                    # 如果这是最后一个响应块且包含完整内容，保存聊天记录
                    if data.get("finished") and account_no:
                        # 获取完整内容并保存
                        full_content = data.get("content", "")
                        if full_content:
                            save_ok = await agent.chat_service.save_chat_message(
                                account_no=account_no,
                                user_msg=query,
                                assistant_msg=full_content
                            )
                            if save_ok:
                                logger.info(f"[account_no={account_no}] 流式聊天记录已保存")
                            # 异步生成对话摘要
                            asyncio.create_task(agent.chat_service.generate_summary(account_no=account_no))
                except json.JSONDecodeError:
                    # 包装非JSON响应
                    wrapped_data = {
                    "content": chunk_str,
                    "chat_id": current_chat_id,
                    "account_no": account_no
                }
                    yield json.dumps(wrapped_data).encode('utf-8')
    except Exception as e:
        error_data = json.dumps({
            "error": f"处理流式响应时出错: {str(e)}",
            "chat_id": current_chat_id,
            "account_no": account_no
        })
        yield error_data.encode('utf-8')